# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: helpers\address\firewall.py
__author__ = "Attila Gerendi (Sanyi)"
import socket, time, logging, ipaddress
from threading import Lock
from helpers.address import is_ip
import yaml

_logger = logging.getLogger("helpers.address.firewall")


class FirewallException(Exception):
    pass


class Firewall(object):
    def __init__(
        self,
        rules,
        reload_interval=None,
        resolve_forbidden_hosts=True,
        forbid_private=True,
    ):
        self._rules = dict()
        self.reload_interval = reload_interval
        self.resolve_forbidden_hosts = resolve_forbidden_hosts
        self._lock = Lock()
        self.input_rules = rules
        self.forbid_private = forbid_private

    @property
    def rules(self):
        with self._lock:
            rules = self._rules
        current_time = time.time()
        if (
            not rules
            or self.reload_interval
            and current_time - rules["last_refresh"] > self.reload_interval
        ):
            pass
        if isinstance(self.input_rules, str):
            try:
                with open(self.input_rules, "r") as (f):
                    _rules = yaml.load(f)
            except Exception as e:
                if not rules:
                    raise
                _logger.exception(
                    "could not load rules file %s: %s", self.input_rules, e
                )
                return rules

            _rules = _rules.get("rules")
            if not _rules:
                if not rules:
                    raise Exception("no rules found")
                _logger.warning("no rules found in %s", self.input_rules)
                return rules
        else:
            if isinstance(self.input_rules, dict):
                _rules = self.input_rules
            else:
                if callable(self.input_rules):
                    _rules = self.input_rules()
                else:
                    raise Exception("invalid firewall rules")
                rules = dict(
                    last_refresh=current_time,
                    ips=set(),
                    hosts=set(_rules.get("hosts", [])),
                    domains=set(_rules.get("domains", [])),
                    ranges=set(),
                )
                ranges = set(_rules.get("ranges", []))
                for cidr_address in ranges:
                    rules["ranges"].add(
                        ipaddress.ip_network(cidr_address, strict=False)
                    )

                ips = set(_rules.get("ips", []))
                for ip in ips:
                    rules["ips"].add(ipaddress.ip_address(ip))

                hosts = rules.get("hosts")
                if self.resolve_forbidden_hosts:
                    if hosts:
                        for host in hosts:
                            try:
                                name, alias_list, address_list = socket.gethostbyname_ex(
                                    host
                                )
                                for ip in address_list:
                                    rules["ips"].add(ipaddress.ip_address(ip))

                            except Exception as e:
                                _logger.exception(
                                    "failed to resolve hosts %s, %s", host, e
                                )

                with self._lock:
                    self._rules = rules
            return rules

    def check_ip(self, ip_address, do_reverse_dns=False):
        rules = self.rules
        if (
            ip_address.is_reserved
            or ip_address.is_link_local
            or ip_address.is_multicast
        ):
            raise FirewallException("forbidden range")
        if self.forbid_private:
            if ip_address.is_private or ip_address.is_loopback:
                raise FirewallException("forbidden range")
        for forbidden_address in rules["ips"]:
            if ip_address == forbidden_address:
                raise FirewallException("forbidden address")

        for forbidden_range in rules["ranges"]:
            if ip_address in forbidden_range:
                raise FirewallException("forbidden range")

        if do_reverse_dns:
            try:
                name, alias_list, _address_list = socket.gethostbyaddr(
                    ip_address.compressed
                )
            except Exception:
                pass
            else:
                self.check_domain(name)
                for alias in alias_list:
                    self.check_domain(alias)

    def check_ip_network(self, cidr_address, max_size=255, do_reverse_dns=False):
        network = ipaddress.ip_network(cidr_address, strict=False)
        if network.prefixlen == 32:
            if network.num_addresses == 1:
                self.check_ip(network.network_address, do_reverse_dns)
        c = 0
        for address in network.hosts():
            c += 1
            if c > max_size:
                raise FirewallException("forbidden address count")
            self.check_ip(address, do_reverse_dns)

        return c

    def check_ip_range(
        self, start_ip_address, stop_ip_address, max_size=256, do_reverse_dns=False
    ):
        networks = ipaddress.summarize_address_range(
            ipaddress.ip_address(start_ip_address),
            ipaddress.ip_address(stop_ip_address),
        )
        c = 0
        try:
            for network in networks:
                for address in network:
                    c += 1
                    if c > max_size:
                        raise FirewallException("forbidden address count")
                    self.check_ip(address, do_reverse_dns)

        except ValueError:
            raise FirewallException("invalid range")

        return c

    def check_domain(self, name):
        rules = self.rules
        name = name.lower()
        if name in rules["hosts"]:
            raise FirewallException("forbidden host")
        for domain in rules["domains"]:
            if name.endswith(domain):
                raise FirewallException("forbidden domain")

    def check_address(self, address, do_reverse_dns=True):
        if is_ip(address):
            self.check_ip(ipaddress.ip_address(address), do_reverse_dns)
        else:
            self.check_domain(address)
            try:
                _, alias_list, ip_address_list = socket.gethostbyname_ex(address)
            except socket.gaierror:
                raise FirewallException("host not found")

            for alias in alias_list:
                self.check_domain(alias)

            for ip in ip_address_list:
                self.check_ip(ipaddress.ip_address(ip), do_reverse_dns)
